# softmax and Cross Entropy
import torch
import torch.nn as nn
import numpy as np


def softmax(x):
    return np.exp(x) / np.sum(np.exp(x), axis=0)


def cross_entropy(actual, predicted):
    loss = -np.sum(actual * np.log(predicted))
    return loss


# x = np.array([2.0, 1.0, 0.1])
# outputs = softmax(x)
# print('softmax numpy: ', outputs)
#
# outputs = torch.softmax(torch.from_numpy(x), dim=0)
# print('torch softmax: ', outputs)
#
# loss = nn.CrossEntropyLoss()

# 3 samples
# Y = torch.tensor([2, 0, 1])
# nsamples x nclasses = 3x3
# Y_predicted_good = torch.tensor([[0.1, 1.0, 2.1], [2.0, 1.0, 0.1], [0.1, 3.0, 0.1]])
# Y_predicted_bad = torch.tensor([[2.1, 1.0, 0.1], [0.1, 1.0, 2.1], [0.1, 3.0, 0.1]])
#
# l1 = loss(Y_predicted_good, Y)
# l2 = loss(Y_predicted_bad, Y)
# print(f'l1: {l1:.4f}, l2: {l2:.4f}')
#
# _, prediction1 = torch.max(Y_predicted_good, 1)
# _, prediction2 = torch.max(Y_predicted_bad, 1)
#
# print(prediction1)
# print(prediction2)

# multi class
class NeuralNets(nn.Module):
    def __init__(self, input_size, hidden_size, num_classes):
        super(NeuralNets, self).__init__()
        self.linear1 = nn.Linear(input_size, hidden_size)
        self.relu = nn.ReLU()
        self.linear2 = nn.Linear(hidden_size, num_classes)

    def forward(self, x):
        out = self.linear1(x)
        out = self.relu(out)
        out = self.linear2(out)
        # no softmax at the end
        return out
model = NeuralNets(input_size=28*28, hidden_size=5, num_classes=3)

criterion = nn.CrossEntropyLoss()

#binary class
class NeuralNet1(nn.Module):
    def __init__(self, input_size, hidden_size):
        super(NeuralNet1, self).__init__()
        self.linear1 = nn.Linear(input_size, hidden_size)
        self.relu = nn.ReLU()
        self.linear2 = nn.Linear(hidden_size, 1)

    def forward(self, x):
        out = self.linear1(x)
        out = self.relu(out)
        out = self.linear2(out)
        # sigmoid at the end
        y_pred = torch.sigmoid(out)
        return y_pred

model = NeuralNet1(input_size=28*28, hidden_size=5)
criterion = nn.BCELoss()